from lib import utils, HSIExtractor, HuInvariantMoment, TextureFeature
import json, sys, os

if len(sys.argv) < 2:
  print('Error: 缺少参数：“要训练的文件夹名”，例：')
  print('python3 train.py ant')
elif not os.path.exists('101_ObjectCategories/' + sys.argv[1]):
  print('Error: 101_ObjectCategories文件夹中没找到' + sys.argv[1])
else:
  category = sys.argv[1]

  data = {
    'color': [],
    'shape': [],
    'texture': []
  }

  for item in sorted(utils.getFilePathListByDirname(category)):
    data['color'].append(HSIExtractor.getFeature(item))
    data['shape'].append(HuInvariantMoment.getFeature(item))
    data['texture'].append(TextureFeature.getFeature(item))

  with open('data/' + category + '.json', 'w') as f:
    f.write(json.dumps(data))
